#!/usr/bin/python
# -*-coding:utf-8-*-

'''
定义一些时间序列技术分析指标算子
'''

import os
import numpy as np
import pandas as pd
from joblib import wrap_non_picklable_objects

# import ta
import talib as talib

from gplearn_evolve.base.utils import d_type_exam
from gplearn_evolve.base.utils import d_map
from gplearn_evolve.base.utils import keys_core_pd
from gplearn_evolve.base.utils import keys_core_pd_v2

# 1
@wrap_non_picklable_objects
def def_keyfunc_ts_ma(keys, base=5, key_by_path=True):
    def gp_ts_ma(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params, min_periods=params//2).mean())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_ma

# 2
@wrap_non_picklable_objects
def def_keyfunc_ts_roc(keys, base=5, key_by_path=True):
    def gp_ts_roc(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.pct_change(params))

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_roc

# 3
@wrap_non_picklable_objects
def def_keyfunc_ts_bias(keys, base=5, key_by_path=True):
    def gp_ts_bias(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:(x / x.rolling(window=params,
                                                                            min_periods=params//2).mean() - 1))

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_bias

# 4
@wrap_non_picklable_objects
def def_keyfunc_ts_ema(keys, base=5, key_by_path=True):
    def gp_ts_ema(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:pd.Series.ewm(x,
                                                                         span=params,
                                                                         min_periods=params-1).mean())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_ema

# 5
@wrap_non_picklable_objects
def def_keyfunc_ts_rsi(keys, base=5, key_by_path=True):
    def gp_ts_rsi(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.astype(np.float64)

        def get_tech_analsis(df, params):
            df = df.copy()

            try:
                df['f1'] = talib.RSI(df['f1'].values, params)
            except:
                pass

            return df[['f1']]

        # values = values.groupby(as_index=False,
        #                         level='stock_code',
        #                         sort=False).apply(lambda x:talib.RSI(x.iloc[:,0], params))
        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(get_tech_analsis, params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        values = values.astype(np.float32)

        del load_keys

        return values['f1'].values

    return gp_ts_rsi

# 6
@wrap_non_picklable_objects
def def_keyfunc_ts_trix(keys, base=5, key_by_path=True):
    def gp_ts_trix(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.astype(np.float64)

        def get_tech_analsis(df, params):
            df = df.copy()

            try:
                df['f1'] = talib.TRIX(df['f1'].values, params)
            except:
                pass

            return df[['f1']]

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(get_tech_analsis, params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        values = values.astype(np.float32)

        del load_keys

        return values['f1'].values

    return gp_ts_trix

# 7
@wrap_non_picklable_objects
def def_keyfunc_ts_kama(keys, base=5, key_by_path=True):
    def gp_ts_kama(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.astype(np.float64)

        def get_tech_analsis(df, params):
            df = df.copy()

            try:
                df['f1'] = talib.KAMA(df['f1'].values, params)
            except:
                pass

            return df[['f1']]

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(get_tech_analsis, params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        values = values.astype(np.float32)

        del load_keys

        return values['f1'].values

    return gp_ts_kama


'''linear regression at time series direction'''
# 8
@wrap_non_picklable_objects
def def_keyfunc_ts_ls(keys, base=10, key_by_path=True):
    def gp_ts_ls(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.astype(np.float64)

        def get_tech_analsis(df, params):
            df = df.copy()

            try:
                df['f1'] = talib.LINEARREG(df['f1'].values, params)
            except:
                pass

            return df[['f1']]

        ls_values = values.groupby(as_index=False,
                                   level='stock_code',
                                   sort=False).apply(get_tech_analsis, params)

        # residual
        values = values - ls_values

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        values = values.astype(np.float32)

        del load_keys

        return values['f1'].values

    return gp_ts_ls

# 9
@wrap_non_picklable_objects
def def_keyfunc_ts_ls_slope(keys, base=10, key_by_path=True):
    def gp_ts_ls_slope(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.astype(np.float64)

        def get_tech_analsis(df, params):
            df = df.copy()

            try:
                df['f1'] = talib.LINEARREG_SLOPE(df['f1'].values, params)
            except:
                pass

            return df[['f1']]

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(get_tech_analsis, params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        values = values.astype(np.float32)

        del load_keys

        return values['f1'].values

    return gp_ts_ls_slope

# 10
@wrap_non_picklable_objects
def def_keyfunc_ts_ls_angle(keys, base=10, key_by_path=True):
    def gp_ts_ls_angle(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.astype(np.float64)

        def get_tech_analsis(df, params):
            df = df.copy()

            try:
                df['f1'] = talib.LINEARREG_ANGLE(df['f1'].values, params)
            except:
                pass

            return df[['f1']]

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(get_tech_analsis, params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        values = values.astype(np.float32)

        del load_keys

        return values['f1'].values

    return gp_ts_ls_angle

# 11
@wrap_non_picklable_objects
def def_keyfunc_ts_ls_intercept(keys, base=10, key_by_path=True):
    def gp_ts_ls_intercept(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.astype(np.float64)

        def get_tech_analsis(df, params):
            df = df.copy()

            try:
                df['f1'] = talib.LINEARREG_INTERCEPT(df['f1'].values, params)
            except:
                pass

            return df[['f1']]

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(get_tech_analsis, params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        values = values.astype(np.float32)

        del load_keys

        return values['f1'].values

    return gp_ts_ls_intercept



